/*
 * simple_set.h
 *
 *  Created on: Jun 3, 2011
 *      Author: tqlong
 */

#ifndef SIMPLE_SET_H_
#define SIMPLE_SET_H_

#include <armadillo>
#include <algorithm>

/** x_i >= 0, sum x_i = 1 */
class CanonicalSimplex
{
public:
  typedef arma::vec domain_type;
  static const bool has_projection_ = true;
public:
  domain_type
  projection(const domain_type& y) const
  {
    domain_type s = sort(y, 1);
    //std::cout << "y=\n" << y << "ysort=\n" << s << std::endl;
    double sum = 0.0, tmax = 0.0;
    bool ok = false;
    for (int i = 0; i < length(y) - 1; i++)
      {
        sum += s[i];
        tmax = (sum - 1) / (i + 1);
        //std::cout << "tmax=" << tmax << " ynext=" << s[i + 1] << std::endl;
        if (tmax >= s[i + 1])
          {
            ok = true;
            break;
          }
      }
    if (!ok) tmax = (sum + s[length(y) - 1] - 1) / length(y);

    //std::cout << "tmax=" << tmax << std::endl;
    for (int i = 0; i < length(y); i++)
      {
        s[i] = std::max(0.0, y[i] - tmax);
      }
    //std::cout << "proj=\n" << s << std::endl;
    return s;
  }

  bool
  indicator(const domain_type& y) const
  {
    double s = sum(y);
    if (fabs(s - 1) > 1e-10) return false;
    for (int i = 0; i < length(y); i++)
      if (y[i] < 0) return false;
    return true;
  }

  int
  length(const domain_type& y) const
  {
    return y.n_elem;
  }
};

/** x_i >= 0, sum x_i <= 1 */
class StandardSimplex
{
public:
  typedef arma::vec domain_type;
  static const bool has_projection_ = true;
public:
  domain_type
  projection(const domain_type& y) const
  {
    domain_type s1 = CanonicalSimplex().projection(y);
    domain_type s2 = y;
    for (int i = 0; i < length(y); i++)
      if (s2[i] < 0) s2[i] = 0;
    if (sum(s2) > 1) return s1;

    double d1 = arma::norm(y - s1, 2), d2 = arma::norm(y - s2, 2);
    return d1 < d2 ? s1 : s2;
  }

  bool
  indicator(const domain_type& y) const
  {
    double s = sum(y);
    if (s > 1) return false;
    for (int i = 0; i < length(y); i++)
      if (y[i] < 0) return false;
    return true;
  }
protected:
  int
  length(const domain_type& y) const
  {
    return y.n_elem;
  }
};

#endif /* SIMPLE_SET_H_ */
